
module FokkerEquation (toFokkerMatrix, fullSolveFM, fullSolveReg) where

import Data.Packed.Matrix
import Data.Packed.Vector
import Numeric.LinearAlgebra.LAPACK
import Tikhonov

type DVector = Vector Double

(!>) :: Fractional a=>Element a=>Vector a->Int->a
v !> i   | i<0      = 0
         | i>=dim v = 0
         | otherwise= v @> i

prepareA :: Fractional a=> Element a=> Vector a->Matrix a
prepareA f0 = buildMatrix (dim f0) (dim f0) bf
        where bf (i, j)   | i==j        = -( (f0 !> (i+1)) - (f0 !> (i-1)) )/2
                          | i==j-1      = (f0 !> i)/2
                          | i==j+1      = - (f0 !> i)/2
                          | otherwise   = 0 

prepareB :: Fractional a=>Element a=> Vector a->Matrix a
prepareB f0 = buildMatrix (dim f0) (dim f0) bf
        where bf (i, j)   | i==j          = (f0 !> (i-1)) - 4*(f0 !> i) + (f0 !> (i+1))
                          | i==j-1        = ((f0 !> (i-1)) + 2*(f0 !> i) - (f0 !> (i+1)))/2
                          | i==j+1        = ((f0 !> (i+1)) + 2*(f0 !> i) - (f0 !> (i-1)))/2
                          | otherwise     = 0


toFokkerMatrix :: (Vector Double, Vector Double, Vector Double)-> (Matrix Double, Matrix Double)
toFokkerMatrix (f0, f1, f2) = (fromBlocks [prepareBlocks f0, prepareBlocks f1], trans $ fromRows [diff f0 f1 f2])
        where prepareBlocks t = [(prepareA t), (prepareB t)]
              diff f0 f1 f2   = join [(zipVectorWith (-) f1 f0), (zipVectorWith (-) f2 f1)]
              
solveFM (a, b) = linearSolveSVDR Nothing a b

fullSolveFM (f0, f1, f2) = splitRow $ solveFM $ toFokkerMatrix (f0, f1, f2)
        where splitRow m = toBlocksEvery numRows 1 m
              numRows    = dim f0

solveReg alpha (a, b) = linsolveReg alpha a b
              
fullSolveReg alpha (f0, f1, f2) = splitRow $ solveReg alpha $ toFokkerMatrix (f0, f1, f2)
        where splitRow m = toBlocksEvery numRows 1 m
              numRows    = dim f0